package com.cat.chapter039;

import java.util.ArrayList;
import java.util.TreeMap;

/**
 * @author cat
 * @description https://leetcode.cn/problems/number-of-atoms/description/
 * @create 2025/7/9 19:48
 * @since JDK17
 */

public class Solution03 {
    int where = 0;
    public String countOfAtoms(String formula) {
        TreeMap<String, Integer> ans = f(formula.toCharArray(), 0);
        return toAns(ans);
    }

    public String toAns(TreeMap<String, Integer> ans) {
        StringBuilder res = new StringBuilder();
        for (String s : ans.keySet()) {
            res.append(s);
            Integer cnt = ans.get(s);
            if (cnt > 1) {
                res.append(cnt);
            }
        }
        return res.toString();
    }

    public void fill(String s, TreeMap<String, Integer> ans, TreeMap<String, Integer> atoms, int num) {
        if (s.length() != 0 || atoms != null) {
            num = num == 0 ? 1 : num;
            if (s.length() != 0) {
                ans.put(s, ans.getOrDefault(s, 0) + num);
            } else {
                for (String string : atoms.keySet()) {
                    ans.put(string, ans.getOrDefault(string, 0) + atoms.get(string) * num);
                }
            }
        }
    }

    public TreeMap<String, Integer> f(char[] s, int cur) {
        int num = 0;
        TreeMap<String, Integer> ans = new TreeMap<>();
        TreeMap<String, Integer> atoms = null;
        StringBuilder sb = new StringBuilder();
        while (cur < s.length && s[cur] != ')') {
            if (s[cur] >= 'A' && s[cur] <= 'Z' || s[cur] == '(') {    // 碰到大写
                fill(sb.toString(), ans, atoms, num);
                atoms = null;
                sb.setLength(0);
                num = 0;
                if (s[cur] == '(') {
                    atoms = f(s, cur + 1);
                    cur = where + 1;
                } else {    // 大写字母
                    sb.append(s[cur++]);
                }
            } else if (s[cur] >= 'a' && s[cur] <= 'z') {
                sb.append(s[cur++]);
            } else {
                num = num * 10 + s[cur++] - '0';
            }
        }
        fill(sb.toString(), ans, atoms, num);
        where = cur;
        return ans;
    }

    public static void main(String[] args) {
        Solution03 solution03 = new Solution03();
        System.out.println(solution03.countOfAtoms("K4(ON(SO3)2)2"));
    }

}
